"""
@author: Viet Nguyen <nhviet1009@gmail.com>
"""

import torch
from lib.env import create_train_env
from lib.model import PPO
import torch.nn.functional as F
from collections import deque
from gym_super_mario_bros.actions import SIMPLE_MOVEMENT, COMPLEX_MOVEMENT, RIGHT_ONLY

num_states = 3
torch.manual_seed(123)
# if opt.action_type == "right":
# actions = RIGHT_ONLY
# elif opt.action_type == "simple":
# actions = SIMPLE_MOVEMENT
# else:
actions = COMPLEX_MOVEMENT
num_actions = len(actions)
world = 1
stage = 1
max_actions = 200
num_global_steps = 5e6

env = create_train_env(world, stage, actions)
obs, info = env.reset()
done = True
curr_step = 0
actions = deque(maxlen=max_actions)
while True:
    curr_step += 1
    action = env.action_space.sample()
    state, reward, done, trunc, info = env.step(action=action)
    done = done or trunc

    # Uncomment following lines if you want to save model whenever level is completed
    # if info["flag_get"]:
    #     print("Finished")
    #     torch.save(local_model.state_dict(),
    #                "{}/ppo_super_mario_bros_{}_{}_{}".format(opt.saved_path, opt.world, opt.stage, curr_step))

    if reward != 0:
        print(f"Reward: {reward}, action: {action}, info: {info}")

    env.render()
    actions.append(action)
    if curr_step > num_global_steps or actions.count(actions[0]) == actions.maxlen:
        done = True
    if done:
        curr_step = 0
        actions.clear()
        state, _ = env.reset()
    state = torch.from_numpy(state)
    if torch.cuda.is_available():
        state = state.cuda()